package org.dromara.common.sse.core

import io.github.oshai.kotlinlogging.KotlinLogging
import org.dromara.common.redis.utils.RedisUtils.publish
import org.dromara.common.redis.utils.RedisUtils.subscribe
import org.dromara.common.sse.dto.SseMessageDto
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter
import java.io.IOException
import java.util.concurrent.ConcurrentHashMap
import java.util.function.Consumer

/**
 * 管理 Server-Sent Events (SSE) 连接
 *
 * @author Lion Li
 * @updater LikeYouDo
 */
open class SseEmitterManager {
    companion object {
        private val log = KotlinLogging.logger { }

        /**
         * 订阅的频道
         */
        private const val SSE_TOPIC: String = "global:sse"

        private val USER_TOKEN_EMITTERS: MutableMap<Long, MutableMap<String, SseEmitter>> = ConcurrentHashMap()

    }

    /**
     * 建立与指定用户的 SSE 连接
     *
     * @param userId 用户的唯一标识符，用于区分不同用户的连接
     * @param token  用户的唯一令牌，用于识别具体的连接
     * @return 返回一个 SseEmitter 实例，客户端可以通过该实例接收 SSE 事件
     */
    fun connect(userId: Long, token: String): SseEmitter {
        // 从 USER_TOKEN_EMITTERS 中获取或创建当前用户的 SseEmitter 映射表（ConcurrentHashMap）
        // 每个用户可以有多个 SSE 连接，通过 token 进行区分
        val emitters: MutableMap<String, SseEmitter> =
            USER_TOKEN_EMITTERS.computeIfAbsent(userId) { ConcurrentHashMap<String, SseEmitter>() }

        // 创建一个新的 SseEmitter 实例，超时时间设置为一天 避免连接之后直接关闭浏览器导致连接停滞
        val emitter = SseEmitter(86400000L)

        emitters[token] = emitter

        // 当 emitter 完成、超时或发生错误时，从映射表中移除对应的 token
        emitter.onCompletion {
            val remove = emitters.remove(token)
            remove?.complete()
        }
        emitter.onTimeout {
            val remove = emitters.remove(token)
            remove?.complete()
        }
        emitter.onError {
            val remove = emitters.remove(token)
            remove?.complete()
        }

        try {
            // 向客户端发送一条连接成功的事件
            emitter.send(SseEmitter.event().comment("connected"))
        } catch (e: IOException) {
            // 如果发送消息失败，则从映射表中移除 emitter
            emitters.remove(token)
        }
        return emitter
    }

    /**
     * 断开指定用户的 SSE 连接
     *
     * @param userId 用户的唯一标识符，用于区分不同用户的连接
     * @param token  用户的唯一令牌，用于识别具体的连接
     */
    fun disconnect(userId: Long?, token: String?) {
        if (userId == null || token == null) {
            return
        }
        val emitters: MutableMap<String, SseEmitter>? = USER_TOKEN_EMITTERS[userId]
        if (!emitters.isNullOrEmpty()) {
            try {
                val sseEmitter: SseEmitter = emitters[token]!!
                sseEmitter.send(SseEmitter.event().comment("disconnected"))
                sseEmitter.complete()
            } catch (ignore: Exception) {
            }
            emitters.remove(token)
        } else {
            USER_TOKEN_EMITTERS.remove(userId)
        }
    }

    /**
     * 订阅SSE消息主题，并提供一个消费者函数来处理接收到的消息
     *
     * @param consumer 处理SSE消息的消费者函数
     */
    fun subscribeMessage(consumer: Consumer<SseMessageDto>?) {
        subscribe(
            SSE_TOPIC,
            SseMessageDto::class.java,
            consumer
        )
    }

    /**
     * 向指定的用户会话发送消息
     *
     * @param userId  要发送消息的用户id
     * @param message 要发送的消息内容
     */
    fun sendMessage(userId: Long, message: String) {
        val emitters: MutableMap<String, SseEmitter>? = USER_TOKEN_EMITTERS[userId]
        if (!emitters.isNullOrEmpty()) {
            for ((key, value) in emitters) {
                try {
                    value.send(
                        SseEmitter.event()
                            .name("message")
                            .data(message)
                    )
                } catch (e: Exception) {
                    val remove = emitters.remove(key)
                    remove?.complete()
                }
            }
        } else {
            USER_TOKEN_EMITTERS.remove(userId)
        }
    }

    /**
     * 本机全用户会话发送消息
     *
     * @param message 要发送的消息内容
     */
    fun sendMessage(message: String) {
        for (userId in USER_TOKEN_EMITTERS.keys) {
            sendMessage(userId, message)
        }
    }

    /**
     * 发布SSE订阅消息
     *
     * @param sseMessageDto 要发布的SSE消息对象
     */
    fun publishMessage(sseMessageDto: SseMessageDto) {
        val broadcastMessage = SseMessageDto().apply {
            message = sseMessageDto.message
            userIds = sseMessageDto.userIds
        }
        publish(SSE_TOPIC, broadcastMessage) {
            log.info { "SSE发送主题订阅消息topic:${SSE_TOPIC} session keys:${sseMessageDto.userIds} message:${sseMessageDto.message}" }
        }
    }

    /**
     * 向所有的用户发布订阅的消息(群发)
     *
     * @param message 要发布的消息内容
     */
    fun publishAll(message: String) {
        val broadcastMessage = SseMessageDto().apply {
            this.message = message
        }
        publish(SSE_TOPIC, broadcastMessage) {
            log.info { "SSE发送主题订阅消息topic:${SSE_TOPIC} message:${message}" }
        }
    }

}
